#!/bin/bash

. cmd.sh
. path.sh

echo "$0 $@"

# config
H=`pwd`
n=24
stage=1
end=5

gpu_mode=true

. utils/parse_options.sh || exit 1;

if [ $# -lt 1 ] || [ $# -gt 5 ]; then
    echo "Usage: $0 [options] <data-name> <data-dir> <dict-path> <out-dir> [<log-dir>]";
    echo "e.g.: $0 BBSpeech data/BBSpeech data/BBSpeech/dict  ali-out/BBSpeech exp/BBSpeech"
    echo "Note: <log-dir> defaults to exp/<data-name>"
    echo "Options: "
    echo "  --nj <nj>                                        # number of parallel jobs"
    echo "  --gpu_mode <mode>                                # [true/wait]"
    exit 1;
fi

data_name=$1
data_path=$2
dict_path=$3
outdir=$4
if [ $# -ge 5 ]; then
    logdir=$5
else
    logdir=exp/${data_name}
fi

DATA=data/$data_name/
mkdir -p $DATA/train_mfcc
mkdir -p $DATA/train_fbank
mkdir -p $DATA/dict
mkdir -p $outdir
mkdir -p $logdir

for f in $data_path/wav.scp $data_path/text $data_path/utt2spk;do
    [ ! -f $f ] && echo "$0: no such file $f" && exit 1;
    cp $f $DATA/train_mfcc
    cp $f $DATA/train_fbank
done

if [ $stage -le 1 ] && [ $end -ge 1 ];then
    perl utils/utt2spk_to_spk2utt.pl $DATA/train_mfcc/utt2spk > $DATA/train_mfcc/spk2utt
    cp $DATA/train_mfcc/spk2utt $DATA/train_fbank/spk2utt
    utils/fix_data_dir.sh $DATA/train_mfcc
    steps/make_mfcc.sh --nj $n --cmd "$train_cmd" $DATA/train_mfcc || exit 1;
    steps/compute_cmvn_stats.sh $DATA/train_mfcc || exit 1;

    steps/make_fbank.sh --nj $n --cmd "$train_cmd" $DATA/train_fbank || exit 1;
    steps/compute_cmvn_stats.sh $DATA/train_fbank || exit 1;
fi

if [ $stage -le 2 ] && [ $end -ge 2 ];then
    [ ! -f $dict_path ] && echo "$0: no dict file $f" && exit 1;
    cp $dict_path data/$data_name/dict
    python3 tools/prep_dict.py $DATA/dict/lexicon.txt
    utils/prepare_lang.sh --position-dependent-phones false $DATA/dict "#" $DATA/lang/tmp $DATA/lang
fi

if [ $stage -le 3 ] && [ $end -ge 3 ];then
    #monophone
    steps/train_mono.sh --boost-silence 1.25 --nj $n --cmd "$train_cmd" $DATA/train_mfcc $DATA/lang $logdir/mono || exit 1;

    #monophone_ali
    steps/align_si.sh --boost-silence 1.25 --nj $n --cmd "$train_cmd" $DATA/train_mfcc $DATA/lang $logdir/mono $logdir/mono_ali || exit 1;

    #triphone
    steps/train_deltas.sh --boost-silence 1.25 --cmd "$train_cmd" 10000 80000 $DATA/train_mfcc $DATA/lang $logdir/mono_ali $logdir/tri1 || exit 1;

    #triphone_ali
    steps/align_si.sh --nj $n --cmd "$train_cmd" $DATA/train_mfcc $DATA/lang $logdir/tri1 $logdir/tri1_ali || exit 1;

    #lda_mllt
    steps/train_lda_mllt.sh --cmd "$train_cmd" --splice-opts "--left-context=3 --right-context=3" 12000 120000 $DATA/train_mfcc $DATA/lang $logdir/tri1_ali $logdir/tri2b || exit 1;

    #lda_mllt_ali
    steps/align_si.sh  --nj $n --cmd "$train_cmd" --use-graphs true $DATA/train_mfcc $DATA/lang $logdir/tri2b $logdir/tri2b_ali || exit 1;

    #sat
    steps/train_sat.sh --cmd "$train_cmd" 15000 180000 $DATA/train_mfcc $DATA/lang $logdir/tri2b_ali $logdir/tri3b || exit 1;

    #sat_ali
    steps/align_fmllr.sh --nj $n --cmd "$train_cmd" $DATA/train_mfcc $DATA/lang $logdir/tri3b $logdir/tri3b_ali || exit 1;

    #quick
    steps/train_quick.sh --cmd "$train_cmd" 18000 250000 $DATA/train_mfcc $DATA/lang $logdir/tri3b_ali $logdir/tri4b || exit 1;

    #quick_ali
    steps/align_fmllr.sh --nj $n --cmd "$train_cmd" $DATA/train_mfcc $DATA/lang $logdir/tri4b $logdir/tri4b_ali || exit 1;
fi

if [ $stage -le 4 ] && [ $end -ge 4 ];then
    local/train_chain.sh --nj $n \
        --gpu_mode $gpu_mode \
        --Exp $logdir \
        --train_mfcc $DATA/train_mfcc \
        --train_fbank $DATA/train_fbank \
        --ali_dir $logdir/tri4b_ali \
        --tri4b $logdir/tri4b \
        --lang $logdir/chain/lang \
        --lang_gmm $DATA/lang \
        --treedir $logdir/chain/tri4_tree \
        $logdir/chain/tdnn_attend
fi

if [ $stage -le 5 ] && [ $end -ge 5 ];then
    # modified align_lats.sh:line 90, fixed frame_subsampling_opt="--frame-subsampling-factor=1".
    steps/nnet3/align_lats.sh --nj $n --cmd "$train_cmd" $DATA/train_fbank $logdir/chain/lang $logdir/chain/tdnn_attend $logdir/chain_ali || exit 1;
    
    # --frame-shift(s) here must match to the value(ms) in conf.
    steps/get_train_ctm.sh --frame-shift 0.004 $DATA/train_fbank $logdir/chain/lang $logdir/chain_ali $outdir/chain_ali

    python3 local/textgrid/ctm_to_textgrid.py $outdir/chain_ali/ctm $outdir/chain_ali/textgrid
    sort -k1 -u $outdir/chain_ali/textgrid.scp -o $outdir/chain_ali/textgrid.scp

    if [ -f $data_path/textgrid.scp ];then
        src_textgrid=$data_path/textgrid.scp
        python3 local/textgrid/count_diff_textgrid.py $data_name $src_textgrid $outdir/chain_ali/textgrid.scp $outdir/chain_ali/
    fi
fi